import copy
import os
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt


class DynamicModel:

    def __init__(
        self,
        Iw: float,
        Ipy: float,
        Ipz: float,
        M: float,
        m: float,
        r: float,
        l: float,
        d: float,
    ):
        self.Iw = Iw
        self.Ipy = Ipy
        self.Ipz = Ipz
        self.M = M
        self.m = m
        self.r = r
        self.l = l
        self.d = d

        # self.x = 0
        # self.dx = 0
        # self.theta = 0
        # self.dtheta = 0
        # self.phi = 0
        # self.dphi = 0
        self.state = np.zeros((6, 1))

        self.matA = np.zeros((6, 6))
        self.matB = np.zeros((6, 2))
        self.updateMatrix()

    def initState(self, state: list[float]):
        self.state = np.zeros((6, 1))
        for i in range(6):
            self.state[i, 0] = state[i]

    def updateMatrix(self):
        """
        A
        0  1  0  0  0  0
        0  0  v1 0  0  0
        0  0  0  1  0  0
        0  0  v3 0  0  0
        0  0  0  0  0  1
        0  0  0  0  0  0
        B
        0  0
        v2 v2
        0  0
        v4 v4
        0  0
        v5 v6
        """
        k1 = self.M + self.m * 2 + 2 * self.Iw / (self.r * self.r)
        A = 1 / (k1 * self.r)
        B = self.M * self.l / k1
        k2 = self.Ipy + self.M * self.l * self.l
        C = -self.M * self.l / k2
        D = 9.8 / k2
        E = 1 / k2
        v1 = -B * D / (1 + B * C)
        v2 = (A + B * E) / (1 + B * C)
        v3 = D / (1 + B * C)
        v4 = (C * A - E) / (1 + B * C)
        v5 = (
            1
            / (
                self.m * self.d
                + self.Iw * self.d / (self.r * self.r)
                + 2 * self.Ipz / self.d
            )
            * self.r
        )
        v6 = -v5
        self.matA = np.array(
            [
                [0, 1, 0, 0, 0, 0],
                [0, 0, v1, 0, 0, 0],
                [0, 0, 0, 1, 0, 0],
                [0, 0, v3, 0, 0, 0],
                [0, 0, 0, 0, 0, 1],
                [0, 0, 0, 0, 0, 0],
            ]
        )
        self.matB = np.array([[0, 0], [v2, v2], [0, 0], [v4, v4], [0, 0], [v5, v6]])
        print("A:")
        print(self.matA)
        print("B:")
        print(self.matB)

    def update(
        self, leftTorque: float, rightTorque: float, dt: float
    ) -> tuple[float, float, float, float]:
        """
        return: speed, gyro, angle, gyroz
        """
        ctrl = np.array([[leftTorque], [rightTorque]])
        # dx = Ax + Bu
        # x(k+1) = x(k)+ dx(k) * dt
        dx = self.matA @ self.state + self.matB @ ctrl
        self.state = self.state + dx * dt
        return self.state[1, 0], self.state[3, 0], self.state[2, 0], self.state[5, 0]


class EVA:
    """
    用于评估动力学模型的准确性能，比较在相同激励下，不同系统模型的输出一致性
    节点订阅控制量并更新记录系统状态，绘制两者差异
    """

    class State:
        def __init__(
            self,
            gyro: float,
            angle: float,
            gyroz: float,
            speed: float,
            tl: float,
            tr: float,
        ):
            self.gyro = gyro
            self.angle = angle
            self.gyroz = gyroz
            self.speed = speed
            self.tl = tl
            self.tr = tr

    def __init__(self, rate: int, model: DynamicModel):
        self.model = model
        self.states1 = []  # 状态记录bag信息
        self.states2 = []  # 状态记录模型信息
        self.state = self.State(0, 0, 0, 0, 0, 0)
        self.rate = rate
        self.df = pd.read_csv("database/record.csv")

    # tl,tr -> x,theta,phi
    def update(self, start: int, step: int):
        self.state = self.State(
            self.df["gyro"][start],
            self.df["angle"][start],
            self.df["gyroz"][start],
            self.df["speed"][start],
            self.df["tl"][start],
            self.df["tr"][start],
        )
        self.model.initState(
            [
                0,
                self.state.speed,
                self.state.angle,
                self.state.gyro,
                0,
                self.state.gyroz,
            ]
        )
        for i in range(step):
            self.state = self.State(
                self.df["gyro"][start + i],
                self.df["angle"][start + i],
                self.df["gyroz"][start + i],
                self.df["speed"][start + i],
                self.df["tl"][start + i],
                self.df["tr"][start + i],
            )
            self.states1.append(copy.copy(self.state))
            speed, gyro, angle, gyroz = self.model.update(
                self.state.tl, self.state.tr, 1.0 / self.rate
            )
            self.states2.append(
                EVA.State(gyro, angle, gyroz, speed, self.state.tl, self.state.tr)
            )
        aix1 = plt.subplot(3, 2, 1)
        aix1.plot([x.gyro for x in self.states1], label="gyro")
        aix1.plot([x.gyro for x in self.states2], label="model")
        aix1.legend()
        aix2 = plt.subplot(3, 2, 2)
        aix2.plot([x.angle for x in self.states1], label="angle")
        aix2.plot([x.angle for x in self.states2], label="model")
        aix2.legend()
        aix3 = plt.subplot(3, 2, 3)
        aix3.plot([x.gyroz for x in self.states1], label="gyroz")
        aix3.plot([x.gyroz for x in self.states2], label="model")
        aix3.legend()
        aix4 = plt.subplot(3, 2, 4)
        aix4.plot([x.speed for x in self.states1], label="speed")
        aix4.plot([x.speed for x in self.states2], label="model")
        aix4.legend()
        aix5 = plt.subplot(3, 2, 5)
        aix5.plot([x.tl for x in self.states1], label="tl")
        aix5.plot([x.tr for x in self.states1], label="tr")
        aix5.legend()
        plt.show()


def main(args=None):
    model = DynamicModel(
        Iw=0.02535,
        Ipy=0.375833,
        Ipz=0.829583,
        M=22,
        m=3,
        r=0.13,
        l=0.05,
        d=0.26 * 2,
    )
    eva_node = EVA(100, model)

    while True:
        start = int(input("start(number)").strip())
        step = int(input("step(number)").strip())
        print("start:", start, "step:", step)
        eva_node.update(start, step)
        input("press enter to continue")


if __name__ == "__main__":
    main()
